# exphandle.py

import logging
import os
import re
import sys

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from tqdm import tqdm


VERSION = "0.4.1"
INFO = f"""
--- *.EXP table tool (ver{VERSION}) ---
This program extracts the exports from the TCm-calc,
developed by FZU-MCMF, main developers: Qiao Yang, Gao Min-Liang, Wu Bo*, et al.
please contact via wubo@fzu.edu.cn, 654489521@qq.com in case.
-----------------------------------
"""


class ExpProcessor:
    def __init__(self, folderpath: str):
        self.folderpath = folderpath
        self.groups = {"Phase Fractions (PFs)": ["XBL-T.EXP"], 
                       "Phase Species Fractions (PSFs)": ["XCF-T.EXP"], 
                       "Sublattice Occupying Fractions (SOFs)": ["YSITE-T.EXP"],
                       "Gibbs energy (G)": ["G-T.EXP", "GI-T.EXP"],
                       "Enthalpy (H)": ["H-T.EXP", "HI-T.EXP"],
                       "Entropy (S)": ["S-T.EXP", "SI-T.EXP"], }

    def _split_exp(self, filepath: str):
        """Split the *-T.EXP file into blocks.

        Yields: list[str]: blocks started with '$ PLOTTED COLUMNS ARE'.
        """
        with open(filepath, "r", encoding="utf-8") as file:
            filelines = file.readlines()
        block = []

        for line in filelines:
            if "$ PLOTTED COLUMNS ARE" in line:
                if len(block) > 0:
                    yield block
                    block = []
            block.append(line)

        # Yield the last block
        if len(block) > 0:
            yield block

    def _extract_data(self, block: list[str]):
        """Extract data from the block.
        Returns: dict[str, list[float]]: data extracted from the block.
        """
        if match := re.search(r": (\S+) and (\S+)", block[0]):
            column1, column2 = "T", match.group(2)
        else:
            return
        column1s = []
        column2s = []
        flag = False
        for line in block[1:]:
            parts = line.split()
            if len(parts) < 2:
                continue
            if "M" == parts[-1]:
                flag = True
            if flag:
                v1, v2 = parts[:2]
                try:
                    column1s.append(float(v1))
                    column2s.append(np.float32(v2))
                except ValueError:
                    continue

        data = {column1: column1s, column2: column2s}
        return data

    def process_exp(self, filepaths: list[str]) -> pd.DataFrame:
        all_data = {}
        for filepath in filepaths:
            if not os.path.exists(filepath):
                logging.warning(f"{filepath} not found.")
                continue

            for block in self._split_exp(filepath):
                data = self._extract_data(block)
                if data is None:
                    continue
                df = pd.DataFrame(data).set_index("T")
                for col in df.columns:
                    if col not in all_data:
                        all_data[col] = df[col]
                    else:
                        all_data[col] = all_data[col].combine_first(df[col])

        if all_data:
            combined_data = pd.concat(all_data, axis=1)
            key = "T"  # combined_data.keys()[0]
            combined_data.reset_index(inplace=True)
            combined_data.sort_values(by=key, ascending=True)
            combined_data.drop_duplicates(subset=key, inplace=True)
            combined_data.set_index(key, inplace=True)
        else:
            combined_data = pd.DataFrame()

        return combined_data

    def plot_df(self, name: str, df: pd.DataFrame, folderpath: str):
        """Plot the dataframe and save the plot as a PNG file"""
        all_markers = ["o", "s", "^", "D", "*", "P", "H", "X",]
        # Find a suitable set of markers based on the number of columns
        num_col = len(df.columns)
        if num_col <= len(all_markers):
            markers = all_markers[:num_col]
        else:
            for i in range(len(all_markers), 0, -1):
                if num_col % i == 0:
                    markers = all_markers * (num_col // i)
                    break
            else:
                markers = all_markers[:num_col]
        num_m = len(markers)
        mark_every = 100  # Add markers every 100 units
        plt.figure(figsize=(12, 8))
        for i, column in enumerate(df.columns):
            plt.plot(df.index, df[column], label=column,
                     marker=markers[i % num_m], markevery=mark_every)

        plt.title(f"{name} vs T", fontsize=20)
        plt.xlabel("T, k", fontsize=20)
        plt.ylabel(name, fontsize=20)
        plt.grid()
        plt.legend(loc="best", bbox_to_anchor=(1, 1), fontsize=10)
        plt.tight_layout(rect=[0, 0, 1, 0.95])
        plt.savefig(os.path.join(folderpath, f"{name}-T.png"))
        plt.close()

    def exp2xls(self, folderpath: str = "", output: str = ""):
        self.folderpath = folderpath or self.folderpath
        dfs: dict[str, pd.DataFrame] = {}
        fs = [f for g in self.groups.values() for f in g]
        pbar = tqdm(total=len(fs), ncols=80, desc=f"Reading {os.path.basename(folderpath)}")
        for name, group in self.groups.items():
            filepaths = [os.path.join(self.folderpath, file) for file in group]
            df = self.process_exp(filepaths)
            if not df.empty:
                dfs[name] = df
            pbar.update(len(group))
        pbar.close()

        if not dfs:
            logging.warning("No data found. Please check the files.")
            return

        output = os.path.join(output or self.folderpath, "tc_exps.xlsx")
        with pd.ExcelWriter(output) as writer:
            for sheet_name, df in tqdm(dfs.items(), ncols=80, desc="Processing"):
                df.to_excel(writer, sheet_name=sheet_name, index=True)
                self.plot_df(sheet_name, df, self.folderpath)

    def process_directory(self, root_directory: str = ""):
        """处理指定根目录下的所有子目录中的EXP文件"""
        root_directory = root_directory or self.folderpath
        
        # 遍历根目录下的所有子目录
        for root, dirs, files in os.walk(root_directory):
            # 检查当前目录是否有EXP文件
            exp_files = [f for f in files if f.upper().endswith('.EXP')]
            if exp_files:
                logging.info(f"Processing directory: {root}")
                try:
                    self.exp2xls(root)
                    logging.info(f"Successfully processed directory: {root}")
                except Exception as e:
                    logging.error(f"Error processing directory {root}: {e}")
            else:
                logging.info(f"No EXP files found in directory: {root}")


def get_folder(folderpath: str):
    while True:
        folderpath = folderpath or input("Enter the folderpath >>> ").strip()
        if os.path.isdir(folderpath):
            return os.path.abspath(folderpath)
        else:
            logging.warning(f"{folderpath} not found. Please try again.")
            folderpath = ""


def main():
    folderpath = sys.argv[1] if len(sys.argv) > 1 else ""

    while True:
        print(INFO)
        try:
            folderpath = get_folder(folderpath)
            exp = ExpProcessor(folderpath)
            for root, dirs, files in os.walk(folderpath):
                if not any(file.upper().endswith(".EXP") for file in files):
                    logging.warning(f"No .EXP files found in {root}")
                    continue

                exp.exp2xls(root)

            folderpath = ""

        except KeyboardInterrupt:
            print("\n")
            logging.info("Thank you for using ExpProcessor! Exiting...")
            sys.exit(0)
        except Exception as e:
            print("\n")
            logging.exception(e)